from typing import List, Optional

from tree_rag.session import Message

from . import DEFAULT_MODEL_NAME, Prompt, client

# 领域外拒答
SYSTEM_PROMPT = """你是拒答子助手，你负责处理领域外的问题。"""


PROMPT = """
请根据以下对话历史和用户的最新输入，判断用户输入的内容是否属于以下类别之一：
- 询问你的身份（例如：你是谁？）: {role_prompt}
- 问候（例如：你好，早上好）
- 感谢（例如：谢谢，太感谢了）
- 不涉及专业知识的日常闲聊（例如：今天天气真好）
{exclusion_block}

如果用户的意图属于上述任意一种情况，请根据你的角色设定和上下文进行自然回复。
如果用户的意图不属于上述任意一种情况，请使用以下固定话术拒答：
{refuse_word}

注意: 不要解释，不要加括号。

对话历史：
{history}

用户最新输入：{user_say}

请你根据情况或按照要求回复，或拒答。
"""


def build_prompt(
    user_say: str,
    histories: List[Message],
    role_prompt: str,
    refuse_word: Optional[str] = None,
    exclusion: Optional[list[dict]] = None,
) -> "Prompt":
    system_prompt = SYSTEM_PROMPT.format()
    history = "\n".join([item.to_str() for item in histories])
    
    # Use provided refuse word or fallback to default
    final_refuse_word = refuse_word or "对不起，我无法提供这方面的信息。如果您有其他需要帮助的地方，欢迎随时提问。"

    # Build exclusion block
    exclusion_lines: List[str] = []
    if exclusion:
        exclusion_lines.append("另外，下列类别也无需使用固定拒答话术，请按类别自然回复：")
        for item in exclusion:
            category = item.get("category") if isinstance(item, dict) else getattr(item, "category", None)
            examples = item.get("examples") if isinstance(item, dict) else getattr(item, "examples", [])
            definition = item.get("definition") if isinstance(item, dict) else getattr(item, "definition", None)
            words = item.get("words") if isinstance(item, dict) else getattr(item, "words", [])
            
            if category:
                # Build a more natural description for each category
                category_desc = f"**{category}**"
                
                if definition:
                    # Clean up definition formatting and add it naturally
                    clean_definition = definition.strip().replace('\n', ' ')
                    category_desc += f"：{clean_definition}"
                
                if examples and len(examples) > 0:
                    examples_str = "、".join(examples)
                    category_desc += f" 例如：{examples_str}。"
                elif definition:
                    category_desc += "。"
                
                if words and len(words) > 0:
                    # Use the first response template as the main one
                    main_response = words[0] if isinstance(words[0], str) else str(words[0])
                    category_desc += f" 建议回复：{main_response}"
                
                exclusion_lines.append(f"- {category_desc}")
    exclusion_block = ("\n" + "\n".join(exclusion_lines)) if exclusion_lines else ""

    user_prompt = PROMPT.format(
        role_prompt=role_prompt,
        history=history, 
        user_say=user_say,
        refuse_word=final_refuse_word,
        exclusion_block=exclusion_block,
    )
    return Prompt(system_prompt=system_prompt, user_prompt=user_prompt)


async def run_agent(
    user_say: str,
    histories: List[Message],
    role_prompt: str,
) -> str:
    """Run the agent.

    Args:
        context: The context.
        query: The query.
    """
    prompt: Prompt = build_prompt(user_say, histories, role_prompt)
    chat_completion = await client.chat.completions.create(
        messages=[
            {"role": "system", "content": prompt.system_prompt},
            {"role": "user", "content": prompt.user_prompt},
        ],
        model=DEFAULT_MODEL_NAME,
    )
    response = chat_completion.choices[0].message.content
    return response
